import os
import frappe
from frappe_dm.database.dm.db_manager import DmDbManager

def setup_database(force, source_sql=None, verbose=False):
    """初始化达梦数据库"""
    frappe.local.session = frappe._dict({"user":"Administrator"})
    
    db_name = frappe.conf.db_name
    root_conn = get_root_connection(frappe.flags.root_login, frappe.flags.root_password)
    dbman = DmDbManager(root_conn)
    print(dbman.get_database_list())
    if force or db_name not in dbman.get_database_list():
        dbman.delete_user(db_name)
        dbman.drop_database(db_name)
    else:
        raise Exception(f"Database {db_name} already exists")

    dbman.create_user(db_name, frappe.conf.db_password)
    if verbose:
        print(f"Created user {db_name}")

    dbman.create_database(db_name)
    if verbose:
        print(f"Created database {db_name}")

    dbman.grant_all_privileges(db_name, db_name)
    if verbose:
        print(f"Granted privileges to user {db_name} and database {db_name}")

    # 关闭root连接
    root_conn.close()

    bootstrap_database(db_name, verbose, source_sql)

def get_root_connection(root_login=None, root_password=None):
    """获取达梦数据库root连接"""
    if not frappe.local.flags.root_connection:
        if not root_login:
            root_login = frappe.conf.get("root_login") or "SYSDBA"

        if not root_password:
            root_password = frappe.conf.get("root_password") or None

        if not root_password:
            from getpass import getpass
            root_password = getpass("达梦数据库管理员密码: ")

        frappe.local.flags.root_connection = frappe.database.get_db(
            host=frappe.conf.db_host,
            port=frappe.conf.db_port,
            user=root_login,
            password=root_password,
        )

    return frappe.local.flags.root_connection

def bootstrap_database(db_name, verbose, source_sql=None):
    """引导安装数据库"""
    import sys
    
    frappe.connect(db_name=db_name)
    
    import_db_from_sql(source_sql, verbose)
    
    frappe.connect(db_name=db_name)
    
    if "tabDefaultValue" not in frappe.db.get_tables():
        print("Database not installed correctly, please check your DM root password and permissions")
        sys.exit(1)

def import_db_from_sql(source_sql=None, verbose=False):
    """从SQL文件导入数据库"""
    if verbose:
        print("Starting database import...")
        
    if not source_sql:
        source_sql = os.path.join(os.path.dirname(__file__), "framework_dm.sql")
        
    DmDbManager.restore_database(
        verbose=verbose,
        source_sql=source_sql
    )
    
    if verbose:
        print(f"Imported from database {source_sql}")

def drop_user_and_database(db_name, root_login, root_password):
    """删除用户和数据库"""
    frappe.local.db = get_root_connection(root_login, root_password)
    dbman = DmDbManager(frappe.local.db)
    dbman.drop_database(db_name)
    dbman.delete_user(db_name)